// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compute/function_internal.h>

#include <nebula/array/util.h>
#include <nebula/compute/function.h>
#include <nebula/compute/registry.h>
#include <nebula/io/memory.h>
#include <nebula/ipc/reader.h>
#include <nebula/ipc/writer.h>
#include <nebula/core/record_batch.h>
#include <nebula/types/scalar.h>
#include <turbo/base/checked_cast.h>

namespace nebula::compute::internal {

    constexpr char kTypeNameField[] = "_type_name";

    turbo::Result<std::shared_ptr<StructScalar>> FunctionOptionsToStructScalar(
            const FunctionOptions &options) {
        std::vector<std::string> field_names;
        std::vector<std::shared_ptr<Scalar>> values;
        const auto *options_type =
                dynamic_cast<const GenericOptionsType *>(options.options_type());
        if (!options_type) {
            return turbo::unimplemented_error("serializing ", options.type_name(),
                                          " to StructScalar");
        }
        TURBO_RETURN_NOT_OK(options_type->ToStructScalar(options, &field_names, &values));
        field_names.push_back(kTypeNameField);
        const char *options_name = options.type_name();
        values.emplace_back(
                new BinaryScalar(Buffer::Wrap(options_name, std::strlen(options_name))));
        return StructScalar::create(std::move(values), std::move(field_names));
    }

    turbo::Result<std::unique_ptr<FunctionOptions>> FunctionOptionsFromStructScalar(
            const StructScalar &scalar) {
        TURBO_MOVE_OR_RAISE(auto type_name_holder, scalar.field(kTypeNameField));
        const std::string type_name =
                turbo::checked_cast<const BinaryScalar &>(*type_name_holder).value->to_string();
        TURBO_MOVE_OR_RAISE(auto raw_options_type,
                              get_function_registry()->get_function_options_type(type_name));
        const auto *options_type = turbo::checked_cast<const GenericOptionsType *>(raw_options_type);
        return options_type->FromStructScalar(scalar);
    }

    turbo::Result<std::shared_ptr<Buffer>> GenericOptionsType::serialize(
            const FunctionOptions &options) const {
        TURBO_MOVE_OR_RAISE(auto scalar, FunctionOptionsToStructScalar(options));
        TURBO_MOVE_OR_RAISE(auto array, MakeArrayFromScalar(*scalar, 1));
        auto batch =
                RecordBatch::create(schema({field("", array->type())}), /*num_rows=*/1, {array});
        TURBO_MOVE_OR_RAISE(auto stream, io::BufferOutputStream::create());
        TURBO_MOVE_OR_RAISE(auto writer, ipc::make_file_writer(stream, batch->schema()));
        TURBO_RETURN_NOT_OK(writer->write_record_batch(*batch));
        TURBO_RETURN_NOT_OK(writer->close());
        return stream->finish();
    }

    turbo::Result<std::unique_ptr<FunctionOptions>> GenericOptionsType::deserialize(
            BufferSpan buffer) const {
        return DeserializeFunctionOptions(buffer);
    }

    turbo::Result<std::unique_ptr<FunctionOptions>> DeserializeFunctionOptions(
            BufferSpan buffer) {
        // Copying the buffer here is not ideal, but we need to do it to avoid
        // use-after-free issues with the zero-copy buffer read.
        auto stream = io::BufferReader::from_string(span_as_string(buffer));
        TURBO_MOVE_OR_RAISE(auto reader, ipc::RecordBatchFileReader::open(stream.get()));
        TURBO_MOVE_OR_RAISE(auto batch, reader->read_record_batch(0));
        if (batch->num_rows() != 1) {
            return turbo::invalid_argument_error(
                    "serialized FunctionOptions's batch repr was not a single row - had ",
                    batch->num_rows());
        }
        if (batch->num_columns() != 1) {
            return turbo::invalid_argument_error(
                    "serialized FunctionOptions's batch repr was not a single column - had ",
                    batch->num_columns());
        }
        auto column = batch->column(0);
        if (column->type()->id() != Type::STRUCT) {
            return turbo::invalid_argument_error(
                    "serialized FunctionOptions's batch repr was not a struct column - was ",
                    column->type()->to_string());
        }
        TURBO_MOVE_OR_RAISE(auto raw_scalar,
                              turbo::checked_cast<const StructArray &>(*column).get_scalar(0));
        auto scalar = turbo::checked_cast<const StructScalar &>(*raw_scalar);
        return FunctionOptionsFromStructScalar(scalar);
    }

    turbo::Status CheckAllArrayOrScalar(const std::vector<Datum> &values) {
        for (const auto &value: values) {
            if (!value.is_value()) {
                return turbo::failed_precondition_error(
                        "Tried executing function with non-array, non-scalar type: ", value.to_string());
            }
        }
        return turbo::OkStatus();
    }

    turbo::Result<std::vector<TypeHolder>> GetFunctionArgumentTypes(const std::vector<Datum> &args) {
        // type-check Datum arguments here. Really we'd like to avoid this as much as
        // possible
        TURBO_RETURN_NOT_OK(CheckAllArrayOrScalar(args));
        std::vector<TypeHolder> inputs(args.size());
        for (size_t i = 0; i != args.size(); ++i) {
            inputs[i] = TypeHolder(args[i].type());
        }
        return inputs;
    }

}  // namespace nebula::compute::internal
